import os
import re
import numpy as np
from osgeo import gdal, ogr, osr
import fiona
from shapely.geometry import shape
from shapely.ops import transform
from tqdm import tqdm
from collections import defaultdict

# ---------------- 配置参数 ----------------
START_DATE = "2024-01-01"
END_DATE = "2024-12-31"
TARGET_CITY = "淄博市"
SHAPEFILE_PATH = "data/shapefile/市.shp"
DOWNLOAD_DIR = "data/modis"
os.makedirs(DOWNLOAD_DIR, exist_ok=True)

# ---------------- 读取 shapefile 并获取目标区域 polygon ----------------
with fiona.open(SHAPEFILE_PATH, 'r') as shapefile:
    target_geom = None
    crs_wkt = None
    for feature in shapefile:
        if feature['properties'].get('市') == TARGET_CITY:
            target_geom = shape(feature['geometry'])
            crs_wkt = shapefile.crs_wkt
            break

    if target_geom is None:
        raise ValueError(f"❌ shapefile 中未找到名为 '{TARGET_CITY}' 的区域")

# ---------------- 过滤 HDF 文件，只保留每日期版本最新的 ----------------


def get_latest_hdf_files(hdf_dir):
    pattern = re.compile(r"A(\d{7})\.h\d{2}v\d{2}\.\d+\.(\d+)\.hdf")
    file_by_date = defaultdict(list)

    for fname in os.listdir(hdf_dir):
        if not fname.endswith(".hdf"):
            continue
        match = pattern.search(fname)
        if match:
            date = match.group(1)
            timestamp = match.group(2)
            file_by_date[date].append((timestamp, fname))
        else:
            file_by_date["unknown"].append(("0", fname))

    hdf_files = []
    for date, files in file_by_date.items():
        latest_file = max(files, key=lambda x: x[0])[1]
        hdf_files.append(latest_file)

    return hdf_files

# ---------------- 读取子数据集 ----------------


def read_subdataset_array(hdf_path, subdataset_suffix):
    print(f"🔎 读取子数据集（后缀匹配）：{subdataset_suffix}")
    hdf = gdal.Open(hdf_path)
    if not hdf:
        raise RuntimeError(f"无法打开 HDF 文件：{hdf_path}")

    subdatasets = hdf.GetSubDatasets()
    matched_name = None
    for name, desc in subdatasets:
        if name.endswith(subdataset_suffix):
            matched_name = name
            break

    if matched_name is None:
        print(f"⚠️ 无此子数据集后缀: {subdataset_suffix}")
        for name, desc in subdatasets:
            print("  ", name)
        return None

    subds = gdal.Open(matched_name)
    return subds.ReadAsArray(), subds.GetGeoTransform(), subds.GetProjectionRef()

# ---------------- 投影 shapefile geometry 到 MODIS 投影 ----------------


def project_geometry(geom, from_crs_wkt, to_wkt):
    from_srs = osr.SpatialReference()
    from_srs.ImportFromWkt(from_crs_wkt)

    to_srs = osr.SpatialReference()
    to_srs.ImportFromWkt(to_wkt)

    transform_func = osr.CoordinateTransformation(from_srs, to_srs)
    return transform(lambda x, y: transform_func.TransformPoint(x, y)[:2], geom)

# ---------------- 主计算逻辑 ----------------


def process_et():
    hdf_files = get_latest_hdf_files(DOWNLOAD_DIR)
    if not hdf_files:
        raise FileNotFoundError("❌ 未找到任何有效的 HDF 文件")

    total_volume = 0.0
    success, fail = 0, 0
    pixel_area = 500 * 500  # 每个像素面积，单位：平方米

    for fname in tqdm(hdf_files, desc="📦 正在处理"):
        path = os.path.join(DOWNLOAD_DIR, fname).replace("\\", "/")
        et_key = 'MOD_Grid_MOD16A2:ET_500m'
        qc_key = 'MOD_Grid_MOD16A2:ET_QC_500m'

        try:
            et_data, gt, proj = read_subdataset_array(path, et_key)
            qc_data, _, _ = read_subdataset_array(path, qc_key)
        except Exception as e:
            print(f"⚠️ 读取失败: {fname} → {e}")
            fail += 1
            continue

        if et_data is None or qc_data is None:
            fail += 1
            continue

        # 投影 shapefile polygon 到 MODIS 投影
        projected_geom = project_geometry(target_geom, crs_wkt, proj)

        # 创建掩膜
        drv = ogr.GetDriverByName('Memory')
        ds = drv.CreateDataSource('mem')
        lyr = ds.CreateLayer('poly', geom_type=ogr.wkbPolygon)
        feat = ogr.Feature(lyr.GetLayerDefn())
        feat.SetGeometry(ogr.CreateGeometryFromWkb(projected_geom.wkb))
        lyr.CreateFeature(feat)

        rows, cols = et_data.shape
        mask_ds = gdal.GetDriverByName('MEM').Create(
            '', cols, rows, 1, gdal.GDT_Byte)
        mask_ds.SetGeoTransform(gt)
        mask_ds.SetProjection(proj)
        gdal.RasterizeLayer(mask_ds, [1], lyr, burn_values=[1])
        mask_array = mask_ds.ReadAsArray().astype(bool)

        # 处理 ET 数据和 QC 掩膜
        et_data = et_data.astype(np.float32)
        et_data[et_data == 32767] = np.nan
        qc_mask = (qc_data == 0) | (qc_data == 1)
        final_mask = mask_array & qc_mask

        et_m = et_data * 0.1 / 1000  # 单位换算为米
        volume = np.nansum(et_m[final_mask]) * pixel_area

        if volume > 0:
            total_volume += volume
            success += 1
        else:
            print(f"⚠️ 蒸散发全部无效: {fname}")
            fail += 1

    print(f"\n✅ 成功处理：{success} 张，❌ 失败：{fail} 张")
    print(f"📊 总蒸散发量为：{total_volume / 1e8:.2f} 亿立方米")


# ---------------- 启动主函数 ----------------
if __name__ == "__main__":
    process_et()
